"""
backbones take 1 input tensor and produce multiple output tensors.

TODO: load corresponding weights (possibly from other frameworks) and 
    save as backbone pretrained .pth files.
"""

import os
import torch

from .vggbase import vggbase

# ------------------------------
# backbones produce 3 tensors
from .vgg import vgg16
from .cspdarknet import cspdarknet53
from .darknet import darknet53, darknet21, load_darknet_weights
from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152


BACKBONE_WEIGHTS = {
    'vgg16': '', 
    'resnet18': '',
    'resnet34': '',
    'resnet50': '',
    'resnet101': '',
    'resnet152': '',
    'darknet53': '',
    'darknet21': '',
    'cspdarknet53': ''
}


def get_backbone_model(name, weights_path=None):
    try:
        model = eval(name)()
    except:
        raise ValueError("backbone model name not supported")
    
    # load weights
    if weights_path is not None and os.path.exists(weights_path):
        state_dict = torch.load(weights_path, map_location='cpu')
        model.load_state_dict(state_dict)
    return model
